import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2

sam_checkpoint = "C:\\Users\\crxc\\PycharmProjects\\segment-anything\\notebooks\\sam_vit_l_0b3195.pth"
model_type = "vit_l"

device = "cuda"

# 生成等间距的点
x = np.linspace(200, 400, 4)
y = np.linspace(200, 400, 4)

# 形成点的网格
xx, yy = np.meshgrid(x, y)
# input_point = np.vstack([xx.ravel(), yy.ravel()]).T
# input_label = np.ones(len(input_point))

input_point = np.array([[300, 300]])
input_label = np.array([1])


# 假设这是你的SAM模型调用函数
def generate_mask_with_sam(predictor):
    masks, scores, logits = predictor.predict(
        point_coords=input_point,
        point_labels=input_label,
        multimask_output=False,
    )
    return masks[0]
    pass


def addPoint(point, label):
    global input_point, input_label
    # 将新点添加到input_point
    input_point = np.append(input_point, [point], axis=0)
    # 将新标签添加到input_label
    input_label = np.append(input_label, label)
    pass


def recall():
    global input_point, input_label
    # 确保数组中至少有两个元素（包括初始值）才能执行撤回操作
    if input_point.shape[0] > 0:
        input_point = input_point[:-1, :]
        input_label = input_label[:-1]
    else:
        print("No points to recall!")
    pass
